import rclpy
from rclpy.node import Node
from nav_msgs.msg import Odometry
from unitree_go.msg import SportModeState,LowState
from tf2_ros import TransformBroadcaster 
from geometry_msgs.msg import TransformStamped
from sensor_msgs.msg import JointState



class Driver(Node):
    def __init__(self):
        super().__init__("driver")
        self.declare_parameter("odom_frame", "odom")
        self.declare_parameter("base_frame", "base")
        self.declare_parameter("publish_tf","true")

        self.odom_frame = self.get_parameter("odom_frame").value
        self.base_frame = self.get_parameter("base_frame").value

        self.odom_pub = self.create_publisher(Odometry, "odom",  10)
        self.mode_sub = self.create_subscription(SportModeState,"/lf/sportmodestate",self.sport_mode_state_callback, 10)
        #  发布tf
        self.tf_broadcaster = TransformBroadcaster(self)
        self.publish_tf = self.get_parameter("publish_tf").value
        #发布四足状态
        self.joint_pub = self.create_publisher(JointState,"joint_states",10)
        # 订阅低层状态
        self.state_sub = self.create_subscription(LowState,"/lf/lowstate",self.low_state_callback,10)
    def low_state_callback(self,state : LowState):
        joint_state = JointState()

        joint_state.header.stamp = self.get_clock().now().to_msg()
        joint_state.name = [
            "FL_hip_joint","FL_thigh_joint","FL_calf_joint",
            "FR_hip_joint","FR_thigh_joint","FR_calf_joint",
            "RL_hip_joint","RL_thigh_joint","RL_calf_joint",
            "RR_hip_joint","RR_thigh_joint","RR_calf_joint"]
        
        for i in range(12):
            
            q = float(state.motor_state[i].q)
            joint_state.position.append(q)

        self.joint_pub.publish(joint_state)

    def sport_mode_state_callback(self, mode :  SportModeState):
        #解析go2的odom并发布AGIROS的odom
        odom = Odometry()
        odom.header.stamp = self.get_clock().now().to_msg()

        odom.header.frame_id = self.odom_frame
        odom.child_frame_id = self.base_frame
        #位置
        odom.pose.pose.position.x = float(mode.position[0])
        odom.pose.pose.position.y = float(mode.position[1])
        odom.pose.pose.position.z = float(mode.position[2])
 
        #线速度
        odom.twist.twist.linear.x = float(mode.velocity[0])
        odom.twist.twist.linear.y = float(mode.velocity[1])
        odom.twist.twist.linear.z = float(mode.velocity[2])

        #角速度
        odom.twist.twist.angular.z = float(mode.yaw_speed)
        #姿态
        odom.pose.pose.orientation.w = float(mode.imu_state.quaternion[0]) #姿态若未变化，则w接近1
        odom.pose.pose.orientation.x = float(mode.imu_state.quaternion[1])
        odom.pose.pose.orientation.y = float(mode.imu_state.quaternion[2])
        odom.pose.pose.orientation.z = float(mode.imu_state.quaternion[3])

        self.odom_pub.publish(odom)

        
        if self.publish_tf:
            trans_form = TransformStamped()
            trans_form.header.stamp = self.get_clock().now().to_msg()

            trans_form.header.frame_id = self.odom_frame
            trans_form.child_frame_id = self.base_frame

            trans_form.transform.translation.x = float(mode.position[0])
            trans_form.transform.translation.y = float(mode.position[1])
            trans_form.transform.translation.z = float(mode.position[2])

            trans_form.transform.rotation = odom.pose.pose.orientation

            self.tf_broadcaster.sendTransform(trans_form)

def main():
    rclpy.init()
    rclpy.spin(Driver())
    rclpy.shutdown()
    
if __name__ == "__main__":
    main()